from peewee import *

from service.das.database_pool import DatabasePool


# 周K、月K、年K数据表定义
class KLineData(Model):
    id = AutoField(primary_key=True)
    stock_code = CharField(index=True)
    trade_date = DateTimeField(index=True)  # 聚合到的时间点
    open = FloatField()
    high = FloatField()
    low = FloatField()
    close = FloatField()
    pre_close = FloatField(null=True)
    change = FloatField(null=True)
    pct_chg = FloatField(null=True)
    vol = FloatField(null=True)
    amount = FloatField(null=True)
    type = CharField(choices=[("W", "周K"), ("M", "月K"), ("Y", "年K")])  # 区分类型

    class Meta:
        database = DatabasePool(max_connections=5).database
        table_name = "kline_data"


class StockKLineDas:
    # 新增K线数据
    def add_kline_data(data):
        """
        新增一条K线数据
        :param data: 字典类型，包含K线数据字段
        """
        KLineData.create(**data)

    # 查询K线数据
    def get_kline_data(stock_code, kline_type, start_date=None, end_date=None):
        """
        查询指定条件的K线数据
        :param stock_code: 股票代码
        :param kline_type: K线类型（'W', 'M', 'Y'）
        :param start_date: 起始日期
        :param end_date: 结束日期
        """
        query = KLineData.select().where(
            KLineData.stock_code == stock_code, KLineData.type == kline_type
        )
        if start_date:
            query = query.where(KLineData.trade_date >= start_date)
        if end_date:
            query = query.where(KLineData.trade_date <= end_date)
        return list(query.dicts())

    # 更新K线数据
    def update_kline_data(kline_id, update_fields):
        """
        更新一条K线数据
        :param kline_id: K线数据的ID
        :param update_fields: 更新的字段和值
        """
        KLineData.update(**update_fields).where(KLineData.id == kline_id).execute()

    # 删除K线数据
    def delete_kline_data(kline_id):
        """
        删除一条K线数据
        :param kline_id: K线数据的ID
        """
        KLineData.delete().where(KLineData.id == kline_id).execute()
